#include <cstdlib>
#include <iostream>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <poll.h>
#include <cstdio>
#include <errno.h>
#include <string>

#include "tcetools/socket.h"

#include "constants.h"
#include "io_loops.h"
#include "halfproxy_auth.h"

using namespace std;
using namespace tcetools;

extern bool debug_info;
char *ctrl_port = NULL;
char *main_port = NULL;
string secret("halfproxy");
int tcp_wndsize = 0;

int main(int argc, char *argv[]) {
	// Parse command line arguements
	if (argc > 0 && (argc % 2)) {
		for (int i = 1; i < argc; i += 2) {
			if (string(argv[i]).compare("-cp") == 0) {
				ctrl_port = argv[i + 1];
			} else if (string(argv[i]).compare("-p") == 0) {
				main_port = argv[i + 1];
			} else if (string(argv[i]).compare("--secret") == 0) {
				secret = argv[i + 1];
			} else if (string(argv[i]).compare("-D") == 0) {
				debug_info = true;
			} else if (string(argv[i]).compare("-w") == 0) {
				tcp_wndsize = atoi(argv[i + 1]);
			}
		}
	}

	if (ctrl_port == NULL || main_port == NULL) {
		cout << "You must specify the control port number and service port number!" << endl;
		cout << "Usage: halfproxy OPTIONS\nList of options:\n"
			 << "\t-cp       <Control port, required>\n"
			 << "\t-p        <Proxy port, required>\n"
			 << "\t--secret  <secret, default: \"halfproxy\">\n"
			 << "\t-w        <TCP window size>\n"
			 << endl;
		return EXIT_FAILURE;
	}

	cout << "Halfproxy bouncer (" << VERSION << ")\n"
		 << "----------------------------"
		 << "\nControl port is: " << ctrl_port
		 << "\nProxy is: " << main_port
		 << "\nShared secret is: " << secret
		 << "\n" << endl;

	// Drop root
	if (geteuid() == 0) { // Hmm.... we are root...
		uid_t real_id = getuid();

		if (real_id == 0) {
			cout << "Current effective user is root, and the real ID is also root. "
					"It is more secure to chown the binary to root, and flip the setuid "
					"bit. Then I can change uid to the real (unprivilaged) "
					"user." << endl;
		} else {
			cout << "Current effective user is root. Changing effective uid to "
					<< real_id << "." << endl;
			seteuid(real_id);
		}
	}

	Socket::ignore_sigpipe();

	while(1) {
		try {
			// Get control connection
			ListenerSocket sbls(ctrl_port, 10);
			cout << "Waiting for control connection... " << flush;
			ServerSocket &ctrl = *(sbls.accept_connection());
			cout << "done" << endl;

			cout << "Starting authentication... " << flush;
			check_auth_server(ctrl, secret, "CONTROL");
			if (ctrl.sck_closed()) {
				cout << "Failed" << endl;
				continue;
			} else cout << "OK" << endl;

			// Wait for incoming connections
			cout << "Opening shop at " << main_port << "..." << endl;
			ListenerSocket ls(main_port, 10);
			struct pollfd pfds[2];
			pfds[0].fd = ls.get_fd();
			pfds[0].events = POLLIN;
			pfds[0].revents = 0;
			
			pfds[1].fd = ctrl.get_fd();
			pfds[1].events = POLLIN;
			pfds[1].revents = 0;

			int pret;

			while ((pret = poll(pfds, 2, 15000)) >= 0) {
				if ((pfds[1].revents & POLLERR) ||
					(pfds[1].revents & POLLHUP) ||
					(pfds[1].revents & POLLNVAL)) {
					cout << "Control connection terminated." << endl;
					break;
				}
				
				if (pfds[1].revents & POLLIN) {
					char buf[2];
					ssize_t readlen = ctrl.raw_read(buf, 2);
					if (readlen == 0) {
						cout << "Control connection terminated." << endl;
						break;
					}
				}
				
				pfds[1].revents = 0;
			
				struct pollfd &pfd = pfds[0];
				
				if ((pfd.revents & POLLERR) ||
					(pfd.revents & POLLHUP) ||
					(pfd.revents & POLLNVAL)) {
					cout << "Listener terminated?!?" << endl;
					return EXIT_FAILURE;
				}

				if (pfd.revents & POLLIN) {
					// Accept incoming connection, fork...
					ServerSocket *client = NULL, *server = NULL;
					pid_t ret = ls.accept_connection_and_fork(client);

					if (ret == -1) {
						cerr << "Warning, unable to fork!" << endl;
					} else if (ret == 0) { // In child process
						cout << "Connection from " << client->get_rem_ip()
							 << "\nRequesting proxy connection... " << flush;
						try {
							// Request connection...
							ctrl.raw_write("GO", 2);
							server = sbls.accept_connection();

							cout << "OK\nAuthenticating... " << flush;
							check_auth_server(*server, secret, "DATA");
							if (server->sck_closed()) {
								cout << "Failed" << endl;
								client->close_sock();
							} else {
								cout << "OK" << endl;
								
								// Start loop
								IOLoop4(client, server, client->get_rem_ip());
							}
						} catch (SocketError &se) {
							// Silently ignore...
							cout << "[Child] Caught a socket error: " << se.what() << endl;
						}

						// Cleanup
						if (client) delete client;
						if (server) delete server;

						// BAI
						_exit(EXIT_SUCCESS);
					} else { // parent process
						// Let's get back them status info to clear them zombie
						// processes
						int status, cleared = 0;
						while (waitpid(-1, &status, WNOHANG) > 0) cleared++;
						if (cleared) cout << "Cleared " << cleared << " zombie(s)." << endl;
					}

				} else if (pret == 0) {
					// Ping... (keepalive)
					ctrl.raw_write("HI", 2);
				}

				pfd.revents = 0;
			}

			if (pret < 0) perror("poll(): ");

			cout << "Closing shop..." << endl;
			ctrl.close_sock();

		} catch (SocketError &se) {
			// Silently ignore...
			cout << "Caught a socket error: " << se.what() << endl;
			sleep(1);
		}
	}
}
